import pandas as pd
import pickle
import numpy as np
from collections import Counter
from io import BytesIO
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings("ignore")

color_dict = {
    99: '#FF0000',  # 红色
    1: '#00FF00',  # 绿色
    2: '#0000FF',  # 蓝色
    3: '#FFA500',  # 橙色
    4: '#800080'  # 紫色
}
type_dict = {1: "WALK",
             2: "VEHICLE",
             3: "TRAIN",
             4: "BICYCLE",
             99: "STAY"}



class plot_travel_chain_no_osm:
    def __init__(self, travel_data):
        self.travel_data = travel_data

    def gene_plot_data(self):
        cut_index = list(self.travel_data[(self.travel_data['transport'].diff() != 0)].index) + [len(self.travel_data)]
        draw_data = []

        for i in range(len(cut_index) - 1):
            draw_data.append(self.travel_data.iloc[cut_index[i]:cut_index[i + 1], :])
        return draw_data

    def plot_travel_chain_no_osm(self, max_lim, min_lim):
        x_max, y_max = max_lim
        x_min, y_min = min_lim
        fig, ax = plt.subplots()
        ax.set_xlim(x_min, x_max)
        ax.set_ylim(y_min, y_max)
        for temp_data in self.gene_plot_data():
            temp_type = temp_data.iloc[0, -1]
            temp_color = color_dict[temp_type]
            type_name = type_dict[temp_type]
            plot_array = temp_data[['longitude', 'latitude']].to_numpy()
            if temp_type == 99:
                ax.scatter(plot_array[:, 0], plot_array[:, 1], s=10, label=type_name, c=temp_color,
                           edgecolor=temp_color, zorder=100)
            else:
                ax.plot(plot_array[:, 0], plot_array[:, 1], linewidth=3, label=type_name, c=temp_color, zorder=20)

        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys())
        return fig, ax


#将出行链数据转化成dataframe形式，用于作图
class trip_data:
    def __init__(self, trip_data):
        self.GENERATED_TRIP = trip_data
    def gen_df(self):
        l = len(self.GENERATED_TRIP)-1
        data_time = []
        for i in range(1, l):
            x=self.GENERATED_TRIP[i][0].item()
            data_time.append(x)
        data_type = []
        for i in range(1, l):
            x=int(self.GENERATED_TRIP[i][3].item())
            data_type.append(x)
        data_lon = []
        for i in range(1, l):
            x=self.GENERATED_TRIP[i][1].item()
            data_lon.append(x)
        data_lat = [] 
        for i in range(1, l):
            x=self.GENERATED_TRIP[i][2].item()
            data_lat.append(x)
        data={
            'time': np.array(data_time),
            'longitude': np.array(data_lon),
            'latitude': np.array(data_lat),
            'transport': np.array(data_type)
        }
        df = pd.DataFrame(data)
        return df
